#############################################################################################
#############################################################################################
#
#   Open Source License/Disclaimer, Forecast Systems Laboratory NOAA/OAR/GSD,
#   325 Broadway Boulder, CO 80305
#
#   This software is distributed under the Open Source Definition, which may be
#   found at http://www.opensource.org/.
#   In particular, redistribution and use in source and binary forms, with or
#   without modification, are permitted provided that the following conditions are met:
#
#   - Redistributions of source code must retain this notice, this list of
#     conditions and the following disclaimer.
#
#   - Redistributions in binary form must provide access to this notice, this
#     list of conditions and the following disclaimer, and the underlying source code.
#
#   - All modifications to this software must be clearly documented, and are
#     solely the responsibility of the agent making the modifications.
#
#   - If significant modifications or enhancements are made to this software,
#     the GSD Software Policy Manager (softwaremgr.fsl@noaa.gov) should be notified.
#
#   THIS SOFTWARE AND ITS DOCUMENTATION ARE IN THE PUBLIC DOMAIN AND ARE
#   FURNISHED "AS IS." THE AUTHORS, THE UNITED STATES GOVERNMENT, ITS INSTRUMENTALITIES,
#   OFFICERS, EMPLOYEES, AND AGENTS MAKE NO WARRANTY, EXPRESS OR IMPLIED, AS TO
#   THE USEFULNESS OF THE SOFTWARE AND DOCUMENTATION FOR ANY PURPOSE. THEY ASSUME
#   NO RESPONSIBILITY (1) FOR THE USE OF THE SOFTWARE AND DOCUMENTATION; OR (2) TO PROVIDE
#   TECHNICAL SUPPORT TO USERS.
#############################################################################################
#############################################################################################

###############################################################################
###############################################################################
# This library contains a suite of functions to generate simulations
# of shallow water model based on Alexander MacDonald's 2014 version  
# , also included are the tangent linear and adjoint functions.
###############################################################################
###############################################################################

try:
    import numpy as np
except:
    raise ImportError('The numpy library cannot be found!')
try:
    import ESMF
except:
    raise ImportError('The ESMP library cannot be found!')

# Import generic Model Class from Shallow Water Model
try:
    from ..Model.Model_Class import ModelClass
except:
    raise ImportError('The ModelClass library cannot be found!')

class ShallowWaterModelClass(ModelClass):
    
    def __init__(self,params):
        
        try:
            nx = int(params['nx'])
        except:
            raise KeyError('The key nx cannot be found in params!')
        try:
            ny = int(params['ny'])
        except:
            raise KeyError('The key ny cannot be found in params!')
        try:
            dx = float(params['dx'])
        except:
            raise KeyError('The key dx cannot be found in params!')            
        try:
            dy = float(params['dy'])
        except:
            raise KeyError('The key dy cannot be found in params!')            
        try:
            ctrx = int(params['x_start'])
        except:
            raise KeyError('The key x_start cannot be found in params!')            
        try:
            ctry = int(params['y_start'])
        except:
            raise KeyError('The key y_start cannot be found in params!')            
        try:
            start_time = int(params['start_time'])
        except:
            raise KeyError('The key start_time cannot be found in params!')            
        try:
            stop_time = int(params['stop_time'])
        except:
            raise KeyError('The key stop_time cannot be found in params!')            
        try:
            timestep = int(params['timestep'])
        except:
            raise KeyError('The key timestep cannot be found in params!')            
 
        if 'history_interval' in params:
            history_interval = int(params['history_interval'])
        else:
            history_interval = timestep
            
        if 'hamp' in params:
            hamp = float(params['hamp'])
        else:
            hamp = 10.
        if 'radius' in params:
            radius = float(params['radius'])
        else:
            radius = 120000.
        if 'hh' in params:
            hh = float(params['hh'])
        else:
            hh=1000.
        if 'dc' in params:
            dc = float(params['dc'])
        else:
            dc=10000.
        if 'u_noise' in params:
            u_noise = float(params['u_noise'])
        else:            
            u_noise=0.011
        if 'v_noise' in params:
            v_noise = float(params['v_noise'])
        else:                       
            v_noise=0.011
        if 'h_noise' in params:
            h_noise = float(params['h_noise'])
        else:                               
            h_noise=0.125       
        if 'use_advection' in params:
            advection = int(params['use_advection'])
        else:
            advection = 1
        if 'use_dynamics' in params:
            dynamics = int(params['use_dynamics'])
        else:
            dynamics = 1
        if 'use_diffusion' in params:
            diffusion = int(params['use_diffusion'])
        else:
            diffusion = 1
        if 'time_diff_choice' in params:
            time_diff_choice = params['time_diff_choice']
        else:
            time_diff_choice = 'rk4'
            
        ModelClass.__init__(self,start_time,stop_time,timestep,history_interval,nx=nx,ny=ny,dx=dx,dy=dy,offset=1,time_diff_choice=time_diff_choice)
        
        self.Name = 'Shallow Water 2D Model'
        self.VarDim = 3  # We only have 3 variables : U, V, H
        
        #*******Parameter Block
        self.ctrx = ctrx        # mountain center in x direction
        self.ctry = ctry        # mountain center in y direction
        self.nxp1=self.nx+1	        # boundary box x direction
        self.nyp1=self.ny+1	        # boundary box y direction
	
        #********** Load constants
        # Logical integer variables to determine model composition
        self.advection=advection	        # advection used
        self.dynamics =dynamics	                # dynamics used
        self.diffusion=diffusion	        # diffusion used

        self.g=10.	        # gravity, m/s**2
        self.hamp=hamp	        # amplitude of initial mound, meters
        self.radius=radius	# radius in grid point space
        
        # create Grids
        self.create_grid( )
        self.x = self.grid.get_coords(0, staggerloc=ESMF.StaggerLoc.CORNER)
        self.y = self.grid.get_coords(1, staggerloc=ESMF.StaggerLoc.CORNER)

        #**************  End Case 1.
        self.nts=self.SimTime.size            # Number of time steps
        
        self.dx2=self.dx*2.	# dx times 2
        self.dy2=self.dy*2.	# dy times 2
    
        self.dt=self.timestep		# time step in seconds
        self.hh=hh	        # height of the mean column (hh=H)
        self.cu=-self.g/(self.dx*2.)	# define constant for u
        self.cv=-self.g/(self.dy*2.)	# define constant for v
        self.ch=-hh/(self.dx*2.)	# define constant for h eqn

        self.dc=dc		# diffusion coeficient
        self.ku=dc*self.timestep/self.dx**2	# u diffusion constant
        self.kv=dc*self.timestep/self.dy**2	# v diffusion constant
        self.kh=dc*self.timestep/self.dx**2	# h diffusion constant 

        self.u_noise = u_noise    # u representative error, m/s
        self.v_noise = v_noise    # v representative error, m/s
        self.h_noise = h_noise    # h representative error, m/s
    
        # Declare Model Variables
        self.u_field = self.create_field('u')
        self.v_field = self.create_field('v')
        self.h_field = self.create_field('h')
     
        # Initial Condition
        self.h_field.data[self.offset:-1-self.offset,self.offset:-1-self.offset] = self.hamp * \
             np.exp((-(self.x[self.offset:-1-self.offset,self.offset:-1-self.offset] - self.ctrx) **2 - \
             (self.y[self.offset:-1-self.offset,self.offset:-1-self.offset] - self.ctry)**2)/self.radius**2)
        self.u_field.data.fill(0.)
        self.v_field.data.fill(0.)
        
        # periodic boundaries     
        self.periodic_boundary(self.h_field.data)
        self.periodic_boundary(self.u_field.data)
        self.periodic_boundary(self.v_field.data)
   
        # Put them in numpy array
        self.h = self.h_field.data
        self.u = self.u_field.data
        self.v = self.v_field.data
        
        # Stack all variables into one array
        self.ic = np.vstack((self.u,self.v,self.h))           
        
    # Self description
    def __repr__(self):
        return '*'*50+'\n'+ \
                self.Name +'\n'+ \
                'Using '+self.time_diff_choice+' time differential scheme to propagate\n'+ \
                '*'*50

    def create_grid(self):
        '''
        PRECONDITIONS: ESMPy has been initialized, 'bounds' contains the 
                    number of indices required for the first two 
                    dimensions of a Grid.  'domask' is a boolean value 
                    that gives the option to put a mask on this Grid.\n
        POSTCONDITIONS: An Grid has been created.\n
        RETURN VALUES: \n Grid :: grid \n
        '''
    
        max_index = np.array([self.nx+self.offset,self.ny+self.offset])
    
        staggerLocs = [ESMF.StaggerLoc.CORNER]
        self.grid = ESMF.Grid(max_index, num_peri_dims=0, staggerloc=staggerLocs)
    
        # VM
        #vm = ESMF.ESMP_VMGetGlobal()
        #localPet, petCount = ESMF.ESMP_VMGet(vm)
    
        # get the coordinate pointers and set the coordinates
        [x,y] = [0, 1]
        gridXCorner = self.grid.get_coords(x, staggerloc=ESMF.StaggerLoc.CORNER)
        gridYCorner = self.grid.get_coords(y, staggerloc=ESMF.StaggerLoc.CORNER)
    
        # make an array that holds indices from lower_bounds to upper_bounds
        bnd2indX = np.arange(self.grid.lower_bounds[ESMF.StaggerLoc.CORNER][x],
                             self.grid.upper_bounds[ESMF.StaggerLoc.CORNER][x], 1)
    
        bnd2indY = np.arange(self.grid.lower_bounds[ESMF.StaggerLoc.CORNER][y],
                             self.grid.upper_bounds[ESMF.StaggerLoc.CORNER][y], 1)
    
        for i in xrange(gridXCorner.shape[x]):
            gridXCorner[i, :] = -self.dx + float(bnd2indX[i])*self.dx 
    
        for j in xrange(gridYCorner.shape[y]):
            gridYCorner[:, j] = -self.dy + float(bnd2indY[j])*self.dy
    
        '''
        # use mpi4py to collect values
        try:
            from mpi4py import MPI
        except:
            raise ImportError("mpi4py is not available, cannot compare \
                            global regridding error")
    
        comm = MPI.COMM_WORLD
        rank = comm.Get_rank()
    
        print "PROC: "+str(rank)
        print "grid bounds: "+str(gridXCoord.shape)
        print "and    bounds: "+str(exLB)+str(exUB)
        '''

    def create_field(self, name):
        '''
        PRECONDITIONS: An Grid has been created, and 'name' is a string that
                    will be used to initialize the name of a new Field.\n
        POSTCONDITIONS: An Field has been created.\n
        RETURN VALUES: \n Field :: field \n
        '''
        # defaults to corner staggerloc
        field = ESMF.Field(self.grid, name,staggerloc=ESMF.StaggerLoc.CORNER)
    
        return field
   
    # Construct a diagonal error matrix
    def error_matrix ( self, error ):
        # Define the noise diagnoal matrix, w/o boundaries
        noise = np.eye((3*self.nx*self.ny))
        
        u_noise,v_noise,h_noise = np.split(error,self.VarDim)
        
        for i in range(self.nx*self.ny):
            noise[i][i] = 1./u_noise**2
        for i in range(self.nx*self.ny, 2*self.nx*self.ny):
            noise[i][i] = 1./v_noise**2
        for i in range(2*self.nx*self.ny, 3*self.nx*self.ny):
            noise[i][i] = 1./h_noise**2
            
        return noise
            
    # Initilize pertubation
    def model_init_tl(self):
        self.g_u = np.zeros_like(self.u)
        self.g_v = np.zeros_like(self.v)
        self.g_h = np.zeros_like(self.h)

        # periodic boundaries
        self.periodic_boundary(self.g_h)
        self.periodic_boundary(self.g_u)
        self.periodic_boundary(self.g_v)
                
        # Stack all variables into one array
        self.ic_tl = np.vstack((self.g_u,self.g_v,self.g_h))
    
    # Initilize gradient
    def model_init_ad(self):
        self.a_u = np.zeros_like(self.u)
        self.a_v = np.zeros_like(self.v)
        self.a_h = np.zeros_like(self.h)

        # periodic boundaries
        self.periodic_boundary(self.a_h)
        self.periodic_boundary(self.a_u)
        self.periodic_boundary(self.a_v)
                
        # Stack all variables into one array
        self.ic_ad = np.vstack((self.a_u,self.a_v,self.a_h))
        
    # Shallow Water model -- Add normal noise on IC
    def model_ic_add_noise(self):
        
        [self.u,self.v,self.h] = np.vsplit(self.ic, self.VarDim)
        
        self.g_u = self.u #+ self.u_noise * np.random.randn(self.u.shape[0],self.u.shape[1])
        self.g_v = self.v #+ self.v_noise * np.random.randn(self.v.shape[0],self.v.shape[1])
        self.g_h = self.h + self.h_noise * np.random.randn(self.h.shape[0],self.h.shape[1])
        
        # periodic boundaries
        self.periodic_boundary(self.g_h)
        self.periodic_boundary(self.g_u)
        self.periodic_boundary(self.g_v)
        
        # Stack all variables into one array
        self.ic_perturbation = np.vstack((self.g_u,self.g_v,self.g_h))
                
    # Shallow Water Model Dynamics
    def model_dyn(self, dudt, dvdt, dhdt ):

        dudt[1:self.nx+1,1:self.ny+1] += self.cu * ( self.h[2:self.nx+1+1,1:self.ny+1  ] - self.h[0:self.nx+1-1,1:self.ny+1  ] )
        dvdt[1:self.nx+1,1:self.ny+1] += self.cv * ( self.h[1:self.nx+1  ,2:self.ny+1+1] - self.h[1:self.nx+1  ,0:self.ny+1-1] )
        dhdt[1:self.nx+1,1:self.ny+1] += self.ch * ( self.u[2:self.nx+1+1,1:self.ny+1  ] - self.u[0:self.nx+1-1,1:self.ny+1  ] + \
                              self.v[1:self.nx+1  ,2:self.ny+1+1] - self.v[1:self.nx+1  ,0:self.ny+1-1] )

    # Shallow Water Model Advection
    def model_adv(self,dudt, dvdt, dhdt ):
        
        dudt[1:self.nx+1,1:self.ny+1] += -self.u[1:self.nx+1,1:self.ny+1] * ( self.u[2:self.nx+1+1,1:self.ny+1] - self.u[0:self.nx+1-1,1:self.ny+1] ) / self.dx2 \
                   -self.v[1:self.nx+1,1:self.ny+1] * ( self.u[1:self.nx+1,2:self.ny+1+1] - self.u[1:self.nx+1,0:self.ny+1-1] ) / self.dy2
        dvdt[1:self.nx+1,1:self.ny+1] += -self.u[1:self.nx+1,1:self.ny+1] * ( self.v[2:self.nx+1+1,1:self.ny+1] - self.v[0:self.nx+1-1,1:self.ny+1] ) / self.dx2 \
                   -self.v[1:self.nx+1,1:self.ny+1] * ( self.v[1:self.nx+1,2:self.ny+1+1] - self.v[1:self.nx+1,0:self.ny+1-1] ) / self.dy2
        dhdt[1:self.nx+1,1:self.ny+1] += -self.u[1:self.nx+1,1:self.ny+1] * ( self.h[2:self.nx+1+1,1:self.ny+1] - self.h[0:self.nx+1-1,1:self.ny+1] ) / self.dx2 \
                   -self.v[1:self.nx+1,1:self.ny+1] * ( self.h[1:self.nx+1,2:self.ny+1+1] - self.h[1:self.nx+1,0:self.ny+1-1] ) / self.dy2

    # Shallow Water Model Diffusion
    def model_dif(self, dudt, dvdt, dhdt ):
        dudt[1:self.nx+1,1:self.ny+1] += self.ku * ( -4. * self.u[1:self.nx+1,1:self.ny+1] + self.u[2:self.nx+1+1,1:self.ny+1] + \
                   self.u[0:self.nx+1-1,1:self.ny+1] + self.u[1:self.nx+1,2:self.ny+1+1] + self.u[1:self.nx+1,0:self.ny+1-1] )
        dvdt[1:self.nx+1,1:self.ny+1] += self.kv * ( -4. * self.v[1:self.nx+1,1:self.ny+1] + self.v[2:self.nx+1+1,1:self.ny+1] + \
                   self.v[0:self.nx+1-1,1:self.ny+1] + self.v[1:self.nx+1,2:self.ny+1+1] + self.v[1:self.nx+1,0:self.ny+1-1] )
        dhdt[1:self.nx+1,1:self.ny+1] += self.kh * ( -4. * self.h[1:self.nx+1,1:self.ny+1] + self.h[2:self.nx+1+1,1:self.ny+1] + \
                   self.h[0:self.nx+1-1,1:self.ny+1] + self.h[1:self.nx+1,2:self.ny+1+1] + self.h[1:self.nx+1,0:self.ny+1-1] )

    # Shallow Water Model Tendency
    def model_tendency(self, ic, time) :
    
        [self.u,self.v,self.h] = np.vsplit(ic, self.VarDim)
        
        dudt = np.zeros_like(self.u)
        dvdt = np.zeros_like(self.v)
        dhdt = np.zeros_like(self.h)
        
        #*********** Call dynamics
        if self.dynamics == 1 : 	        # include dynamics
            self.model_dyn(dudt,dvdt,dhdt)
        # dynamics
   	
        #*********** Call advection
        if self.advection == 1 :             # include advection
            self.model_adv(dudt,dvdt,dhdt)
        # advection
    
        #*********** Call diffusion
        if self.diffusion == 1 :             #include diffusion
            self.model_dif(dudt,dvdt,dhdt)
        # diffusion
        
        # periodic boundaries
        self.periodic_boundary(dhdt)
        self.periodic_boundary(dudt)
        self.periodic_boundary(dvdt)
                
        return np.vstack((dudt,dvdt,dhdt))

      
    # Shallow Water Model Dynamics  --- Tangent Linear
    def model_dyn_tl(self, dudt, dg_udt, dvdt, dg_vdt, dhdt, dg_hdt ):
        # Basic States update
        dudt[1:self.nx+1,1:self.ny+1] += self.cu * ( self.h[2:self.nx+1+1,1:self.ny+1  ] - self.h[0:self.nx+1-1,1:self.ny+1  ] )
        dvdt[1:self.nx+1,1:self.ny+1] += self.cv * ( self.h[1:self.nx+1  ,2:self.ny+1+1] - self.h[1:self.nx+1  ,0:self.ny+1-1] )
        dhdt[1:self.nx+1,1:self.ny+1] += self.ch * ( self.u[2:self.nx+1+1,1:self.ny+1  ] - self.u[0:self.nx+1-1,1:self.ny+1  ] + \
                              self.v[1:self.nx+1  ,2:self.ny+1+1] - self.v[1:self.nx+1  ,0:self.ny+1-1])
        # Perturbation States update
        dg_udt[1:self.nx+1,1:self.ny+1] += self.cu * ( self.g_h[2:self.nx+1+1,1:self.ny+1  ] - self.g_h[0:self.nx+1-1,1:self.ny+1  ] )
        dg_vdt[1:self.nx+1,1:self.ny+1] += self.cv * ( self.g_h[1:self.nx+1  ,2:self.ny+1+1] - self.g_h[1:self.nx+1  ,0:self.ny+1-1] )
        dg_hdt[1:self.nx+1,1:self.ny+1] += self.ch * ( self.g_u[2:self.nx+1+1,1:self.ny+1  ] - self.g_u[0:self.nx+1-1,1:self.ny+1  ] + \
                                self.g_v[1:self.nx+1  ,2:self.ny+1+1] - self.g_v[1:self.nx+1  ,0:self.ny+1-1] )

    # Shallow Water Model Advection --- Tangent Linear
    def model_adv_tl(self, dudt, dg_udt, dvdt, dg_vdt, dhdt, dg_hdt):
        # Basic States update
        dudt[1:self.nx+1,1:self.ny+1] += \
                    - self.u[1:self.nx+1,1:self.ny+1] * ( self.u[2:self.nx+1+1,1:self.ny+1  ] - self.u[0:self.nx+1-1,1:self.ny+1  ] ) / self.dx2 \
                    - self.v[1:self.nx+1,1:self.ny+1] * ( self.u[1:self.nx+1  ,2:self.ny+1+1] - self.u[1:self.nx+1  ,0:self.ny+1-1] ) / self.dy2
        dvdt[1:self.nx+1,1:self.ny+1] += \
                    - self.u[1:self.nx+1,1:self.ny+1] * ( self.v[2:self.nx+1+1,1:self.ny+1  ] - self.v[0:self.nx+1-1,1:self.ny+1  ] ) / self.dx2 \
                    - self.v[1:self.nx+1,1:self.ny+1] * ( self.v[1:self.nx+1  ,2:self.ny+1+1] - self.v[1:self.nx+1  ,0:self.ny+1-1] ) / self.dy2
        dhdt[1:self.nx+1,1:self.ny+1] += \
                    - self.u[1:self.nx+1,1:self.ny+1] * ( self.h[2:self.nx+1+1,1:self.ny+1  ] - self.h[0:self.nx+1-1,1:self.ny+1  ] ) / self.dx2 \
                    - self.v[1:self.nx+1,1:self.ny+1] * ( self.h[1:self.nx+1  ,2:self.ny+1+1] - self.h[1:self.nx+1  ,0:self.ny+1-1] ) / self.dy2
        
        # Perturbation States update
        dg_udt[1:self.nx+1,1:self.ny+1] += \
                     - self.u[1:self.nx+1,1:self.ny+1] * ( self.g_u[2:self.nx+1+1,1:self.ny+1  ] - self.g_u[0:self.nx+1-1,1:self.ny+1  ] ) / self.dx2 \
                     - self.v[1:self.nx+1,1:self.ny+1] * ( self.g_u[1:self.nx+1  ,2:self.ny+1+1] - self.g_u[1:self.nx+1  ,0:self.ny+1-1] ) / self.dy2 \
                     - self.g_u[1:self.nx+1,1:self.ny+1] * ( self.u[2:self.nx+1+1,1:self.ny+1  ] - self.u[0:self.nx+1-1,1:self.ny+1  ] ) / self.dx2 \
                     - self.g_v[1:self.nx+1,1:self.ny+1] * ( self.u[1:self.nx+1  ,2:self.ny+1+1] - self.u[1:self.nx+1  ,0:self.ny+1-1] ) / self.dy2
        dg_vdt[1:self.nx+1,1:self.ny+1] += \
                     - self.u[1:self.nx+1,1:self.ny+1] * ( self.g_v[2:self.nx+1+1,1:self.ny+1  ] - self.g_v[0:self.nx+1-1,1:self.ny+1  ] ) / self.dx2 \
                     - self.v[1:self.nx+1,1:self.ny+1] * ( self.g_v[1:self.nx+1  ,2:self.ny+1+1] - self.g_v[1:self.nx+1  ,0:self.ny+1-1] ) / self.dy2 \
                     - self.g_u[1:self.nx+1,1:self.ny+1] * ( self.v[2:self.nx+1+1,1:self.ny+1  ] - self.v[0:self.nx+1-1,1:self.ny+1  ] ) / self.dx2 \
                     - self.g_v[1:self.nx+1,1:self.ny+1] * ( self.v[1:self.nx+1  ,2:self.ny+1+1] - self.v[1:self.nx+1  ,0:self.ny+1-1] ) / self.dy2
        dg_hdt[1:self.nx+1,1:self.ny+1] += \
                     - self.u[1:self.nx+1,1:self.ny+1] * ( self.g_h[2:self.nx+1+1,1:self.ny+1  ] - self.g_h[0:self.nx+1-1,1:self.ny+1  ] ) / self.dx2 \
                     - self.v[1:self.nx+1,1:self.ny+1] * ( self.g_h[1:self.nx+1  ,2:self.ny+1+1] - self.g_h[1:self.nx+1  ,0:self.ny+1-1] ) / self.dy2 \
                     - self.g_u[1:self.nx+1,1:self.ny+1] * ( self.h[2:self.nx+1+1,1:self.ny+1  ] - self.h[0:self.nx+1-1,1:self.ny+1  ] ) / self.dx2 \
                     - self.g_v[1:self.nx+1,1:self.ny+1] * ( self.h[1:self.nx+1  ,2:self.ny+1+1] - self.h[1:self.nx+1  ,0:self.ny+1-1] ) / self.dy2
        
    # Shallow Water Model Diffusion --- Tangent Linear
    def model_dif_tl(self, dudt, dg_udt, dvdt, dg_vdt, dhdt, dg_hdt ):
    
        # Basic States update
        dudt[1:self.nx+1,1:self.ny+1] += self.ku * ( -4. * self.u[1:self.nx+1,1:self.ny+1] + self.u[2:self.nx+1+1,1:self.ny+1] + \
                self.u[0:self.nx+1-1,1:self.ny+1] + self.u[1:self.nx+1,2:self.ny+1+1] + self.u[1:self.nx+1,0:self.ny+1-1] )
        dvdt[1:self.nx+1,1:self.ny+1] += self.kv * ( -4. * self.v[1:self.nx+1,1:self.ny+1] + self.v[2:self.nx+1+1,1:self.ny+1] + \
                self.v[0:self.nx+1-1,1:self.ny+1] + self.v[1:self.nx+1,2:self.ny+1+1] + self.v[1:self.nx+1,0:self.ny+1-1] )
        dhdt[1:self.nx+1,1:self.ny+1] += self.kh * ( -4. * self.h[1:self.nx+1,1:self.ny+1] + self.h[2:self.nx+1+1,1:self.ny+1] + \
                self.h[0:self.nx+1-1,1:self.ny+1] + self.h[1:self.nx+1,2:self.ny+1+1] + self.h[1:self.nx+1,0:self.ny+1-1] )
        # Perturbation States update
        dg_udt[1:self.nx+1,1:self.ny+1] += self.ku * ( -4. * self.g_u[1:self.nx+1,1:self.ny+1] + self.g_u[2:self.nx+1+1,1:self.ny+1] + \
                  self.g_u[0:self.nx+1-1,1:self.ny+1] + self.g_u[1:self.nx+1,2:self.ny+1+1] + self.g_u[1:self.nx+1,0:self.ny+1-1] )
        dg_vdt[1:self.nx+1,1:self.ny+1] += self.kv * ( -4. * self.g_v[1:self.nx+1,1:self.ny+1] + self.g_v[2:self.nx+1+1,1:self.ny+1] + \
                  self.g_v[0:self.nx+1-1,1:self.ny+1] + self.g_v[1:self.nx+1,2:self.ny+1+1] + self.g_v[1:self.nx+1,0:self.ny+1-1] )
        dg_hdt[1:self.nx+1,1:self.ny+1] += self.kh * ( -4. * self.g_h[1:self.nx+1,1:self.ny+1] + self.g_h[2:self.nx+1+1,1:self.ny+1] + \
                  self.g_h[0:self.nx+1-1,1:self.ny+1] + self.g_h[1:self.nx+1,2:self.ny+1+1] + self.g_h[1:self.nx+1,0:self.ny+1-1] ) 
    
    # Shallow Water Model Tendency --- Tangent Linear
    def model_tendency_tl(self, ic, ic_tl, time) :
          
        [self.u,self.v,self.h] = np.vsplit(ic, self.VarDim)
        [self.g_u,self.g_v,self.g_h] = np.vsplit(ic_tl, self.VarDim)
        
        dudt = np.zeros_like(self.g_u)
        dvdt = np.zeros_like(self.g_v)
        dhdt = np.zeros_like(self.g_h)  
        
        dg_udt = np.zeros_like(self.g_u)
        dg_vdt = np.zeros_like(self.g_v)
        dg_hdt = np.zeros_like(self.g_h)  
        
        #*********** Call dynamics
        if self.dynamics == 1 : 	        # include dynamics
            self.model_dyn_tl(dudt,dg_udt,dvdt,dg_vdt,dhdt,dg_hdt)
        # dynamics
   	
        #*********** Call advection
        if self.advection == 1 :             # include advection
            self.model_adv_tl(dudt,dg_udt,dvdt,dg_vdt,dhdt,dg_hdt)
        # advection
    
        #*********** Call diffusion
        if self.diffusion == 1 :             #include diffusion
            self.model_dif_tl(dudt,dg_udt,dvdt,dg_vdt,dhdt,dg_hdt)
        # diffusion

        # periodic boundaries
        self.periodic_boundary(dhdt)
        self.periodic_boundary(dudt)
        self.periodic_boundary(dvdt)
        self.periodic_boundary(dg_hdt)
        self.periodic_boundary(dg_udt)
        self.periodic_boundary(dg_vdt)
                
        return np.vstack((dudt,dvdt,dhdt)), np.vstack((dg_udt,dg_vdt,dg_hdt))
            
        
    #Shallow Water Model Dynamics --- Adjoint Model
    def model_dyn_ad(self, da_udt, da_vdt, da_hdt ):
        
        ub_ip1_j = self.a_u[self.nx+1: 1:-1, self.ny  : 0:-1]
        ub_im1_j = self.a_u[self.nx-1:  :-1, self.ny  : 0:-1]
        vb_i_jp1 = self.a_v[self.nx  : 0:-1, self.ny+1: 1:-1]
        vb_i_jm1 = self.a_v[self.nx  : 0:-1, self.ny-1:  :-1]
        hb_i_jp1 = self.a_h[self.nx  : 0:-1, self.ny+1: 1:-1]
        hb_i_jm1 = self.a_h[self.nx  : 0:-1, self.ny-1:  :-1]
        hb_ip1_j = self.a_h[self.nx+1: 1:-1, self.ny  : 0:-1]
        hb_im1_j = self.a_h[self.nx-1:  :-1, self.ny  : 0:-1]
        
        dudtb_i_j = da_udt[self.nx:0:-1, self.ny:0:-1]
        dvdtb_i_j = da_vdt[self.nx:0:-1, self.ny:0:-1]
        dhdtb_i_j = da_hdt[self.nx:0:-1, self.ny:0:-1]
            
        ub_ip1_j[:] += + self.ch * dhdtb_i_j
        ub_im1_j[:] += - self.ch * dhdtb_i_j
        vb_i_jp1[:] += + self.ch * dhdtb_i_j
        vb_i_jm1[:] += - self.ch * dhdtb_i_j
        #dhdtb_i_j.fill(0.)
        
        hb_i_jp1[:] += + self.cv * dvdtb_i_j
        hb_i_jm1[:] += - self.cv * dvdtb_i_j
        #dvdtb_i_j.fill(0.)
        
        hb_ip1_j[:] += + self.cu * dudtb_i_j
        hb_im1_j[:] += - self.cu * dudtb_i_j
        #dudtb_i_j.fill(0.)
                
    
    # Shallow Water Model Advection --- Adjoint
    def model_adv_ad(self, da_udt, da_vdt, da_hdt):
        
        ub_i_j   = self.a_u[self.nx  : 0:-1, self.ny  : 0:-1]
        ub_ip1_j = self.a_u[self.nx+1: 1:-1, self.ny  : 0:-1]
        ub_im1_j = self.a_u[self.nx-1:  :-1, self.ny  : 0:-1]
        ub_i_jp1 = self.a_u[self.nx  : 0:-1, self.ny+1: 1:-1]
        ub_i_jm1 = self.a_u[self.nx  : 0:-1, self.ny-1:  :-1]
        
        vb_i_j   = self.a_v[self.nx  : 0:-1, self.ny  : 0:-1]
        vb_ip1_j = self.a_v[self.nx+1: 1:-1, self.ny  : 0:-1]
        vb_im1_j = self.a_v[self.nx-1:  :-1, self.ny  : 0:-1]
        vb_i_jp1 = self.a_v[self.nx  : 0:-1, self.ny+1: 1:-1]
        vb_i_jm1 = self.a_v[self.nx  : 0:-1, self.ny-1:  :-1]
        
        hb_i_jp1 = self.a_h[self.nx  : 0:-1, self.ny+1: 1:-1]
        hb_i_jm1 = self.a_h[self.nx  : 0:-1, self.ny-1:  :-1]
        hb_ip1_j = self.a_h[self.nx+1: 1:-1, self.ny  : 0:-1]
        hb_im1_j = self.a_h[self.nx-1:  :-1, self.ny  : 0:-1]
        
        dudtb_i_j = da_udt[self.nx:0:-1, self.ny:0:-1]
        dvdtb_i_j = da_vdt[self.nx:0:-1, self.ny:0:-1]
        dhdtb_i_j = da_hdt[self.nx:0:-1, self.ny:0:-1]
    
        ub_i_j[:]   += - ( self.v[self.nx+1:1:-1, self.ny:0:-1] - self.v[self.nx-1: :-1, self.ny:0:-1] ) * dvdtb_i_j / self.dx2 \
                       - ( self.u[self.nx+1:1:-1, self.ny:0:-1] - self.u[self.nx-1: :-1, self.ny:0:-1] ) * dudtb_i_j / self.dx2 \
                       - ( self.h[self.nx+1:1:-1, self.ny:0:-1] - self.h[self.nx-1: :-1, self.ny:0:-1] ) * dhdtb_i_j / self.dx2
        hb_ip1_j[:] += - ( self.u[self.nx:0:-1, self.ny:0:-1] * dhdtb_i_j / self.dx2 )
        hb_im1_j[:] += + ( self.u[self.nx:0:-1, self.ny:0:-1] * dhdtb_i_j / self.dx2 )
        vb_i_j[:]   += - ( self.h[self.nx:0:-1, self.ny+1:1:-1] - self.h[self.nx:0:-1, self.ny-1: :-1] ) * dhdtb_i_j / self.dy2
        hb_i_jp1[:] += - ( self.v[self.nx:0:-1, self.ny:0:-1] * dhdtb_i_j / self.dy2 )
        hb_i_jm1[:] += + ( self.v[self.nx:0:-1, self.ny:0:-1] * dhdtb_i_j / self.dy2 )
        
        vb_ip1_j[:] += - ( self.u[self.nx:0:-1, self.ny:0:-1] * dvdtb_i_j / self.dx2 )
        vb_im1_j[:] += + ( self.u[self.nx:0:-1, self.ny:0:-1] * dvdtb_i_j / self.dx2 )
        vb_i_j[:]   += - ( self.v[self.nx:0:-1, self.ny+1:1:-1] - self.v[self.nx:0:-1, self.ny-1: :-1] ) * dvdtb_i_j / self.dy2
        vb_i_jp1[:] += - ( self.v[self.nx:0:-1, self.ny:0:-1] * dvdtb_i_j / self.dy2 )
        vb_i_jm1[:] += + ( self.v[self.nx:0:-1, self.ny:0:-1] * dvdtb_i_j / self.dy2 )
        
        ub_ip1_j[:] += - ( self.u[self.nx:0:-1, self.ny:0:-1] * dudtb_i_j / self.dx2 )
        ub_im1_j[:] += + ( self.u[self.nx:0:-1, self.ny:0:-1] * dudtb_i_j / self.dx2 )
        vb_i_j[:]   += - ( self.u[self.nx:0:-1, self.ny+1:1:-1] - self.u[self.nx:0:-1, self.ny-1: :-1] ) * dudtb_i_j / self.dy2
        ub_i_jp1[:] += - ( self.v[self.nx:0:-1, self.ny:0:-1] * dudtb_i_j / self.dy2 )
        ub_i_jm1[:] += + ( self.v[self.nx:0:-1, self.ny:0:-1] * dudtb_i_j / self.dy2 )

        
    # Shallow Water Model Diffusion --- Adjoint
    def model_dif_ad(self, da_udt, da_vdt, da_hdt ):
    
        ub_i_j    = self.a_u[self.nx  : 0:-1, self.ny  : 0:-1]
        vb_i_j    = self.a_v[self.nx  : 0:-1, self.ny  : 0:-1]
        hb_i_j    = self.a_h[self.nx  : 0:-1, self.ny  : 0:-1]
        
        ub_ip1_j = self.a_u[self.nx+1: 1:-1, self.ny  : 0:-1]
        ub_im1_j = self.a_u[self.nx-1:  :-1, self.ny  : 0:-1]
        ub_i_jp1 = self.a_u[self.nx  : 0:-1, self.ny+1: 1:-1]
        ub_i_jm1 = self.a_u[self.nx  : 0:-1, self.ny-1:  :-1]
        
        vb_ip1_j = self.a_v[self.nx+1: 1:-1, self.ny  : 0:-1]
        vb_im1_j = self.a_v[self.nx-1:  :-1, self.ny  : 0:-1]
        vb_i_jp1 = self.a_v[self.nx  : 0:-1, self.ny+1: 1:-1]
        vb_i_jm1 = self.a_v[self.nx  : 0:-1, self.ny-1:  :-1]
        
        hb_i_jp1 = self.a_h[self.nx  : 0:-1, self.ny+1: 1:-1]
        hb_i_jm1 = self.a_h[self.nx  : 0:-1, self.ny-1:  :-1]
        hb_ip1_j = self.a_h[self.nx+1: 1:-1, self.ny  : 0:-1]
        hb_im1_j = self.a_h[self.nx-1:  :-1, self.ny  : 0:-1]
        
        dudtb_i_j = da_udt[self.nx:0:-1, self.ny:0:-1]
        dvdtb_i_j = da_vdt[self.nx:0:-1, self.ny:0:-1]
        dhdtb_i_j = da_hdt[self.nx:0:-1, self.ny:0:-1]
        
        hb_ip1_j += + self.kh * dhdtb_i_j
        hb_i_j   += - 4. * self.kh * dhdtb_i_j
        hb_im1_j += + self.kh * dhdtb_i_j
        hb_i_jp1 += + self.kh * dhdtb_i_j
        hb_i_jm1 += + self.kh * dhdtb_i_j
        
        vb_ip1_j += + self.kv * dvdtb_i_j
        vb_i_j   += - 4. * self.kv * dvdtb_i_j
        vb_im1_j += + self.kv * dvdtb_i_j
        vb_i_jp1 += + self.kv * dvdtb_i_j
        vb_i_jm1 += + self.kv * dvdtb_i_j
        
        ub_ip1_j += + self.ku * dudtb_i_j
        ub_i_j   += - 4. * self.ku * dudtb_i_j
        ub_im1_j += + self.ku * dudtb_i_j
        ub_i_jp1 += + self.ku * dudtb_i_j
        ub_i_jm1 += + self.ku * dudtb_i_j      
    
    
    # Shallow Water Model Tendency --- Adjoint
    def model_tendency_ad(self, basic_states, forcing, time) :
    
        # Restore the basic states at current time
        [self.u, self.v, self.h] = np.vsplit(basic_states, self.VarDim)
        
        # Define the adjoint states of tendency
        da_udt = np.zeros_like(self.a_u)
        da_vdt = np.zeros_like(self.a_v)
        da_hdt = np.zeros_like(self.a_h)  
        
        self.a_u.fill(0.)
        self.a_v.fill(0.)
        self.a_h.fill(0.)
        
        [forcing_u, forcing_v, forcing_h] = np.vsplit(forcing, self.VarDim)
        da_udt = da_udt + forcing_u
        da_vdt = da_vdt + forcing_v
        da_hdt = da_hdt + forcing_h

        # periodic boundaries adjoint
        self.periodic_boundary_ad(da_hdt)
        self.periodic_boundary_ad(da_udt)
        self.periodic_boundary_ad(da_vdt)
                
        #*********** Call diffusion
        if self.diffusion == 1 :             #include diffusion
            self.model_dif_ad(da_udt,da_vdt,da_hdt)
        # diffusion
   	
        #*********** Call advection
        if self.advection == 1 :             # include advection
            self.model_adv_ad(da_udt,da_vdt,da_hdt)
        # advection
   	
        #*********** Call dynamics
        if self.dynamics == 1 : 	        # include dynamics
            self.model_dyn_ad(da_udt,da_vdt,da_hdt)
        # dynamics
 			
        da_udt.fill(0.)
        da_vdt.fill(0.)   
        da_hdt.fill(0.)
        
        return np.vstack((self.a_u,self.a_v,self.a_h))

    # Update periodic boundary
    def periodic_boundary(self, field):
        
        field[ 0,:] = field[-2,:]
        field[-1,:] = field[ 1,:]       
        field[:, 0] = field[:,-2]
        field[:,-1] = field[:, 1]

    # Update periodic boundary
    def periodic_boundary_ad(self, field):
        
        field[-2,:] += field[ 0,:]
        field[ 0,:] = 0.0
        field[ 1,:] += field[-1,:]  
        field[-1,:] = 0.0
        field[:,-2] += field[:, 0]
        field[:, 0] = 0.0
        field[:, 1] += field[:,-1]
        field[:,-1] = 0.0
    
    # Strip off the boundries
    def strip_off_boundaries(self, fields):
        
        clean_fields = []
        for field in np.split(fields, self.VarDim):
            tmp = np.zeros((self.nx,self.ny))
            tmp[:,:] = field[self.offset:-self.offset, self.offset:-self.offset]
            clean_fields.append( tmp )
         
        return np.vstack(clean_fields)
    
    # patch the boundries
    def patch_boundaries(self, fields):
        
        expanded_fields = []
        for field, field_template in zip( np.split(fields, self.VarDim), np.split(self.ic, self.VarDim)) :
            new_field = np.empty_like(field_template)
            new_field[self.offset:-self.offset, self.offset:-self.offset] = field.reshape((self.nx,self.ny))
            self.periodic_boundary(new_field)
            expanded_fields.append( new_field )
         
        return np.vstack(expanded_fields)